package com.tuniu.agents.common;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.util.DefaultIndenter;
import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.github.victools.jsonschema.generator.Option;
import com.github.victools.jsonschema.generator.SchemaGenerator;
import com.github.victools.jsonschema.generator.SchemaGeneratorConfig;
import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder;
import com.github.victools.jsonschema.module.jackson.JacksonModule;
import com.github.victools.jsonschema.module.jackson.JacksonOption;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.converter.StructuredOutputConverter;
import org.springframework.ai.util.JacksonUtils;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.lang.NonNull;

import java.lang.reflect.Type;
import java.util.Map;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static org.springframework.ai.util.LoggingMarkers.SENSITIVE_DATA_MARKER;

/**
 * An implementation of {@link org.springframework.ai.converter.StructuredOutputConverter} that transforms the LLM output
 * to a specific object type using JSON schema. This converter works by generating a JSON
 * schema based on a given Java class or parameterized type reference, which is then used
 * to validate and transform the LLM output into the desired type.
 *
 * @param <T> The target type to which the output will be converted.
 * @author Mark Pollack
 * @author Christian Tzolov
 * @author Sebastian Ullrich
 * @author Kirk Lund
 * @author Josh Long
 * @author Sebastien Deleuze
 * @author Soby Chacko
 * @author Thomas Vitale
 */
public class CustomerBeanOutputConverter<T> implements StructuredOutputConverter<T> {

    private final Logger logger = LoggerFactory.getLogger(org.springframework.ai.converter.BeanOutputConverter.class);

    /**
     * The target class type reference to which the output will be converted.
     */
    private final Type type;

    /** The object mapper used for deserialization and other JSON operations. */
    private final ObjectMapper objectMapper;

    /** Holds the generated JSON schema for the target type. */
    private String jsonSchema;

    /**
     * Constructor to initialize with the target type's class.
     * @param clazz The target type's class.
     */
    public CustomerBeanOutputConverter(Class<T> clazz) {
        this(ParameterizedTypeReference.forType(clazz));
    }

    /**
     * Constructor to initialize with the target type's class, a custom object mapper, and
     * a line endings normalizer to ensure consistent line endings on any platform.
     * @param clazz The target type's class.
     * @param objectMapper Custom object mapper for JSON operations.
     */
    public CustomerBeanOutputConverter(Class<T> clazz, ObjectMapper objectMapper) {
        this(ParameterizedTypeReference.forType(clazz), objectMapper);
    }

    /**
     * Constructor to initialize with the target class type reference.
     * @param typeRef The target class type reference.
     */
    public CustomerBeanOutputConverter(ParameterizedTypeReference<T> typeRef) {
        this(typeRef.getType(), null);
    }

    /**
     * Constructor to initialize with the target class type reference, a custom object
     * mapper, and a line endings normalizer to ensure consistent line endings on any
     * platform.
     * @param typeRef The target class type reference.
     * @param objectMapper Custom object mapper for JSON operations.
     */
    public CustomerBeanOutputConverter(ParameterizedTypeReference<T> typeRef, ObjectMapper objectMapper) {
        this(typeRef.getType(), objectMapper);
    }

    /**
     * Constructor to initialize with the target class type reference, a custom object
     * mapper, and a line endings normalizer to ensure consistent line endings on any
     * platform.
     * @param type The target class type.
     * @param objectMapper Custom object mapper for JSON operations.
     */
    private CustomerBeanOutputConverter(Type type, ObjectMapper objectMapper) {
        Objects.requireNonNull(type, "Type cannot be null;");
        this.type = type;
        this.objectMapper = objectMapper != null ? objectMapper : getObjectMapper();
        generateSchema();
    }

    /**
     * Generates the JSON schema for the target type.
     */
    private void generateSchema() {
        JacksonModule jacksonModule = new JacksonModule(JacksonOption.RESPECT_JSONPROPERTY_REQUIRED,
                JacksonOption.RESPECT_JSONPROPERTY_ORDER);
        SchemaGeneratorConfigBuilder configBuilder = new SchemaGeneratorConfigBuilder(
                com.github.victools.jsonschema.generator.SchemaVersion.DRAFT_2020_12,
                com.github.victools.jsonschema.generator.OptionPreset.PLAIN_JSON)
                .with(jacksonModule)
                .with(Option.FORBIDDEN_ADDITIONAL_PROPERTIES_BY_DEFAULT);
        SchemaGeneratorConfig config = configBuilder.build();
        SchemaGenerator generator = new SchemaGenerator(config);
        JsonNode jsonNode = generator.generateSchema(this.type);
        ObjectWriter objectWriter = this.objectMapper.writer(new DefaultPrettyPrinter()
                .withObjectIndenter(new DefaultIndenter().withLinefeed(System.lineSeparator())));
        try {
            this.jsonSchema = objectWriter.writeValueAsString(jsonNode);
        }
        catch (JsonProcessingException e) {
            logger.error("Could not pretty print json schema for jsonNode: {}", jsonNode);
            throw new RuntimeException("Could not pretty print json schema for " + this.type, e);
        }
    }

    /**
     * Parses the given text to transform it to the desired target type.
     * @param text The LLM output in string format.
     * @return The parsed output in the desired target type.
     */
    @SuppressWarnings("unchecked")
    @Override
    public T convert(@NonNull String text) {
        try {
            // Remove leading and trailing whitespace
            text = text.trim();

            // Check for and remove triple backticks and "json" identifier
            if (text.startsWith("```") && text.endsWith("```")) {
                // Remove the first line if it contains "```json"
                String[] lines = text.split("\n", 2);
                if (lines[0].trim().equalsIgnoreCase("```json")) {
                    text = lines.length > 1 ? lines[1] : "";
                }
                else {
                    text = text.substring(3); // Remove leading ```
                }

                // Remove trailing ```
                text = text.substring(0, text.length() - 3);

                // Trim again to remove any potential whitespace
                text = text.trim();
            }
            return (T) this.objectMapper.readValue(text, this.objectMapper.constructType(this.type));
        }
        catch (JsonProcessingException e) {
            logger.error(SENSITIVE_DATA_MARKER,
                    "Could not parse the given text to the desired target type: \"{}\" into {}", text, this.type);
            // If parsing fails, try to determine the string result and parse again
            try {
                // Remove leading and trailing whitespace
                text = text.trim();

                // Look for JSON content between ```json and ```
                Pattern pattern = Pattern.compile("```(?:json)?\\s*([\\s\\S]*?)\\s*```", Pattern.CASE_INSENSITIVE);
                Matcher matcher = pattern.matcher(text);

                if (matcher.find()) {
                    // Extract the matched JSON content
                    text = matcher.group(1).trim();
                } else {
                    // If no code block format is found, try to parse the raw text directly
                    // May need to find the starting position of the JSON content
                    int jsonStart = text.indexOf('{');
                    int arrayStart = text.indexOf('[');

                    if (jsonStart >= 0 || arrayStart >= 0) {
                        int startPos = (jsonStart >= 0 && arrayStart >= 0)
                                ? Math.min(jsonStart, arrayStart)
                                : Math.max(jsonStart, arrayStart);
                        text = text.substring(startPos).trim();
                    }
                }

                return (T) this.objectMapper.readValue(text, this.objectMapper.constructType(this.type));
            }
            catch (JsonProcessingException newE) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * Configures and returns an object mapper for JSON operations.
     * @return Configured object mapper.
     */
    protected ObjectMapper getObjectMapper() {
        return JsonMapper.builder()
                .addModules(JacksonUtils.instantiateAvailableModules())
                .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
                .build();
    }

    /**
     * Provides the expected format of the response, instructing that it should adhere to
     * the generated JSON schema.
     * @return The instruction format string.
     */
    @Override
    public String getFormat() {
        String template = """
				Your response should be in JSON format.
				Do not include any explanations, only provide a RFC8259 compliant JSON response following this format without deviation.
				Do not include markdown code blocks in your response.
				Remove the ```json markdown from the output.
				Here is the JSON Schema instance your output must adhere to:
				```%s```
				""";
        return String.format(template, this.jsonSchema);
    }

    /**
     * Provides the generated JSON schema for the target type.
     * @return The generated JSON schema.
     */
    public String getJsonSchema() {
        return this.jsonSchema;
    }

    public Map<String, Object> getJsonSchemaMap() {
        try {
            return this.objectMapper.readValue(this.jsonSchema, Map.class);
        }
        catch (JsonProcessingException ex) {
            logger.error("Could not parse the JSON Schema to a Map object", ex);
            throw new IllegalStateException(ex);
        }
    }

}
